#!/usr/bin/python
#-*- coding: utf-8 -*-
import random, binascii
__author__ = "lyricconch@gmail.com"
__version__ = "1.0b"

PRIMES=[2]  
for i in xrange(3,10000,2):
    for p in PRIMES:
        if i%p==0: break;
        if i/p>=p: continue;
        PRIMES.append(i); break;
del i, p

ISAPRIME=[0,0,0]
def isaprime(aprime, ROUND=(100,20)):
    '''Prime test with miller-rabin method, p=1-.25**R1
    @params aprime: the number to be test
    @params ROUND: R0=rounds of "mod" test, R1=rounds of "m-b" method
    @return bool: True when prime'''   
    ISAPRIME[0]+=1;
    for prime in PRIMES[:ROUND[0]]: 
        if aprime%prime==0: return aprime==prime;    
    if aprime<PRIMES[ROUND[0]-1]**2: return True;
    ISAPRIME[1]+=1;
    dprime=aprime-1; odd=dprime; rank=0;
    while odd&1==0: odd>>=1; rank+=1; 
    for prime in PRIMES[:ROUND[1]] :
        mod= pow(prime, odd, aprime) 
        if mod==1 or mod==dprime: continue;
        for _ in xrange(rank):
            mod= mod**2%aprime
            if mod==dprime: break;
            if mod==1: return False;
        else: return False;
    ISAPRIME[2]+=1;
    return True 
 
def getprime(bits=0, BOUND=[3,1000]):
    """Generate a prime with given bits or bound.
    @params bits: the given bits, when set to 0, use BOUND instead
    @params BOUND: cache the last bound.
    @return long: a prime"""
    if bits:
        bits= max(int(bits),5)
        BOUND[0]=inf= 2**(bits-1)+1
        BOUND[1]=sup= 1<<bits
    else: inf,sup=BOUND;
    for retry in xrange(1000000):
        aprime= random.randrange(inf, sup, 2)   
        if isaprime(aprime): break;
        if retry%100==10: random.seed();    
    print "GETPRIME(%d): call ISAPRIME %d times"%(bits,retry+1)
    return aprime       
   
def generate(bits, adjust=0.9, factor=0.9):
    """Generate keys for rsa crypt/decrypt.
    @params bits: strength of crypt. deteminds Nb=bits
    @params adjuest: to reduce encrypt/keygen. Eb=adjuest*Nb
    @params factor: to prevent factor N. Pb=factor*Nb/2    
    @return (P,Q,N,M,E,D,Nb). bits=>N, adjust=>E, factor=>PQ."""
    P= getprime(bits*factor/2.0+1)
    Q= getprime(bits*(1.9999-factor)/2.0)
    while P==Q: Q=getprime();    
    N=P*Q; M=(P-1)*(Q-1);
    Nb= bitlen(N)
    E= getprime(Nb*adjust)
    while E>=M or gcd(M,E)!=1: E=getprime()
    D= ordmod(M, E)    
    return N,E,D,P,Q,Nb   
    
def translate(data, baseN, keyED, DEC=False, STATIC={}):    
    """RSA main function 
    @params data: data to be encrypt/decrypt
    @params baseN: RSA public param N
    @params keyED: RSA encrypt/decrypt key E/D
    @params DEC: action perform, decrypt when true
    @params STATIC: cache something 
    @return (data,pad) data encrypted/decrypted, len of pad byte"""
    try: temp=data; size=STATIC[baseN];
    except KeyError:         
        size=STATIC[baseN]= -(bitlen(baseN)/-8)    
        #print "! Nbit=%d, Craw/Crsa=%d/%d, Cost=%.3f%%"%(
        #    bitlen(baseN),size-1,size,100.0/(size-1))           
    if not DEC: 
        pad= (-2-len(data))%(size-1) 
        temp= temp+" "*pad+chr(pad/256)+chr(pad%256)  
        sobj2=size*2; ssrc2=sobj2-2;
    else: ssrc2=size*2; sobj2=ssrc2-2;
    assert len(temp)%(size if DEC else size-1)==0    
    temp= binascii.b2a_hex(temp)
    temp= "".join([hex(pow(long(temp[pd:pd+ssrc2],16),keyED,baseN))
        [2:-1].zfill(sobj2) for pd in xrange(0,len(temp),ssrc2)])
    temp= binascii.a2b_hex(temp)              
    assert len(temp)%(size-1 if DEC else size)==0
    if DEC: 
        pad= ord(temp[-2])*256+ord(temp[-1])
        temp= temp[:-2-pad]
    """; print "@ groups=%d lenpad=%d lendata=%d lentran=%d\
        \n  data: %.60s%s\n  tran: %.60s%s"%(
        -(len(temp)/-sobj2), pad, len(data), len(temp),
        binascii.b2a_hex(data), 30<len(data)and"..."or"",
        binascii.b2a_hex(temp), 30<len(temp)and"..."or"") #"""
    return temp, pad   
    
def testRSA(bits=256, length=1000, loops=1, adjust=0.99):  
    from time import clock as now     
    VERBOSE= loops<0; 
    raw= ('__love you more than I can Say! '*1000)[:length];
    if loops<0: 
        loops= -loops 
    t0= now(); 
    N,E,D,P,Q,Nb= generate(bits,adjust,0.8)             
    print 'init: time=%.3f Nb=%d Eb=%d Db=%d Pb=%d Qb=%d'%tuple(
        [now()-t0]+[bitlen(i) for i in (N,E,D,P,Q)])  
    if VERBOSE: 
        if D<2**32: print "  E: %d  D: %d  N: %d  P: %d  Q: %d"%(E,D,N,P,Q);      
        else: print"  E: %XH\n  D: %XH\n  N: %XH\n  P: %XH\n  Q: %XH"%(E,D,N,P,Q);
    t0=now(); temp=raw; lenraw=len(raw);
    for _ in xrange(loops): 
        temp,lenpad= translate(temp,N,E)
    t1=now(); rsa=temp; lenrsa=len(rsa);
    for _ in xrange(loops): 
        temp,lenpad= translate(temp,N,D,1);         
    t2=now(); tz=t2-t0; groups= lenrsa-lenraw-2;
    assert temp==raw #and lenpad==pad    
    print 'info: loops=%d padding=%d groups=%d lenraw=%d lenrsa=%d'%(
        loops, lenpad, groups, lenraw, lenrsa)
    print 'time: .Enc=%.3f .Dec=%.3f /Loop=%.3f /KiB=%.3f /Chunk=%.3f'%(
        t1-t0, t2-t1, tz/loops, 1000*tz/lenraw, tz/loops/groups)    
    if VERBOSE: print "draw: %.60s%s\ndrsa: %.120s%s"%(
        raw, 60<len(raw)and"..."or"", 
        binascii.b2a_hex(rsa), 60<len(rsa)and"..."or"")   
    
def bitlen(num):
    try: return num.bit_length();
    except: 
        if num<0: 
            num=-num   
        bit=0
        while num!=0: 
            num>>=1
            bit+=1
        return bit
        
def ordmod(M, E):
    if E>M: E=E%M;
    assert M%E!=0
    xa=1; xb=0; xc=M;
    ya=0; yb=1; yc=E;
    while yc!=1:
        xc, yc, t= yc, xc%yc, xc/yc        
        xb, yb= yb, xb-yb*t 
        xa, ya= ya, xa-ya*t
    D= yb if yb>0 else M+yb
    assert E*D%M == 1    
    return D  

def gcd(p, q):
    assert p>=1 and q>=1
    while q!=0: p,q=q,p%q;
    else: return p;
    
def factor(n):
    result= []
    for p in PRIMES:
        factor=p; rank=0;            
        while True:
            d,r= divmod(n,factor)
            if r!=0: break;
            n,rank= d,rank+1;
        if rank!=0: 
            result.append((factor,rank))                            
        if n==1 or n<=p*p: break;
    if n!=1: result.append((n,1))
    return result    
    
def phi(n):
    result= n
    for f,r in factor(n):
        result= result/f*(f-1)
    return result

def miu(n):
    result= 1
    for f,r in factor(n):
        if r>1: return 0;
        result= -result
    return result    
    
      
        
if __name__ == '__main__':    
    primes= [40094690950920881030683735292761468389214899724061]
    composites= [1373653]
    print "Primes:"
    for i in primes:
        print 'ISAPRIME(%d): %s' %(i, isaprime(i))    
    print "Composites:"
    for i in composites:
        print 'ISAPRIME(%d): %s' %(i, isaprime(i))  
    
    testRSA(512,10000,-1,0.6);    
    
""" AKS prime test:
    Input: integer n > 1.
    If n = a**b for integers a > 0 and b > 1, output composite.
    Find the smallest r such fits s:(s*n%r==1) > log2(n)**2.
    If 1 < gcd(a,n) < n for some a <= r, output composite.
    If n <= r, output prime.
    For a = 1 to floor(sqrt(phi(r))*log2(n)) do
        if (X+a)**n != X**n+a (mod X**r-1,n), output composite;
    Output prime.

isaprime_AKS(N):
    Nb= N.bit_length()
    gate=Nb**2; r=gate+1;
    while ordmod(r,N)<=gate: r+=1;        
    for i in xrange(3, r+1, 2):
        if gcd(i, N)!=1: return False;
    if N<=r: return True;    
    for i in xrange(1,Nb*sqrt(phi(r))):
        pass
    
translate(data, baseN, keyED, DEC=False, SIZE={}):    
    try: size= SIZE[baseN];
    except KeyError:        
        Nb= bitlen(baseN)
        SIZE[baseN]=size= (Nb+7)/8;
        print "! Nbit=%d, Craw/Crsa=%d/%d, Cost=%.3f%%"%(
            Nb,size-1,size,100.0/(size-1))       
    if not DEC:
        si,so= size-1,2*size;
        pad= -(len(data)+2)%(size-1)
        data+= " "*pad+chr(pad/256)+chr(pad%256)
    else: si,so= size,2*size-2; assert len(data)%si==0;
    tpl="0%dX"%so    
    tran= "".join([a2b_hex(format(pow(long(b2a_hex(data[pd:pd+si]),16),
        keyED,baseN),tpl)) for pd in xrange(0,len(data),si)])   
    if DEC: 
        pad= ord(tran[-2])*256+ord(tran[-1])
        tran= tran[:-pad-2] 
    else: assert len(tran)%size==0;
    print "@ groups=%d lenpad=%d lendata=%d lentran=%d\
        \n  data: %.60s%s\n  tran: %.60s%s"%(
        len(data)/si, pad, len(data), len(tran),
        b2a_hex(data), 30<len(data)and"..."or"",
        b2a_hex(tran), 30<len(tran)and"..."or"") 
    return tran,pad
 """
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
        
""""""